from sklearn.linear_model import Perceptron, LogisticRegression
import numpy as np
from sklearn.metrics import accuracy_score

from 宋海文.机器学习.init import train_dataset, val_dataset, test_dataset


# 将数据转换为 NumPy 数组（Sklearn 需要 2D 输入）
def flatten_data(dataset):
    X = []
    y = []
    for img, label in dataset:
        X.append(img.numpy().flatten())  # 展平成 1D
        y.append(label)
    return np.array(X), np.array(y)

# 处理训练集和测试集
X_train, y_train = flatten_data(train_dataset)
X_val, y_val = flatten_data(val_dataset)
X_test, y_test = flatten_data(test_dataset)

# 训练感知机
perceptron = Perceptron(max_iter=1000, tol=1e-3)
perceptron.fit(X_train, y_train)

# 评估模型
y_pred = perceptron.predict(X_test)
print("Perceptron Test Accuracy:", accuracy_score(y_test, y_pred))


# 训练逻辑回归模型
logistic = LogisticRegression(max_iter=1000)
logistic.fit(X_train, y_train)

# 评估模型
y_pred_logistic = logistic.predict(X_test)
print("Logistic Regression Test Accuracy:", accuracy_score(y_test, y_pred_logistic))
